# -*- coding: UTF-8 -*-

import numpy as np
import pandas as pd
import os
from sklearn.preprocessing import MinMaxScaler
import xlrd


# from statsmodels.tsa.seasonal import seasonal_decompose
# import matplotlib.pylab as plt


def getlen_data(root_dir, data_name, sheetname):
    data_dir = os.path.join(root_dir, data_name + '.xlsx')
    dataframe = pd.read_excel(data_dir, sheet_name=sheetname)
    data_len = len(dataframe)

    return data_len


def load_data(root_dir, data_name, data_len, test_len, sheetname, feature, normal):
    data_dir = os.path.join(root_dir, data_name + '.xlsx')
    dataframe = pd.read_excel(data_dir, sheet_name=sheetname)
    data = dataframe.values[:, 1:]
    data = np.reshape(data, [data.shape[0], feature])
    if normal is True:
        scaler = MinMaxScaler(feature_range=(0, 1))
        data = scaler.fit_transform(data)
    train_data = data[:(data_len - test_len), :]
    test_data = data[(data_len - test_len):, :]

    return train_data, test_data


def generate_data(seq, look_back):
    X = []  # 初始化输入序列X
    Y = []  # 初始化输出序列Y
    for i in range(len(seq) - look_back):
        X.append(seq[i:i + look_back])  # 从输入序列第一期出发，等步长连续不间断采样
        Y.append(seq[i + look_back, 0])  # 对应每个X序列的滞后一期序列值
    return np.array(X, dtype=np.float32), np.array(Y, dtype=np.float32)


def generate_multidata(seq, look_back, look_ahead):
    X = []  # 初始化输入序列X
    Y = []  # 初始化输出序列Y
    for i in range(len(seq) - look_back - look_ahead):
        X.append(seq[i:i + look_back])  # 从输入序列第一期出发，等步长连续不间断采样
        Y.append(seq[i + look_back:i + look_back + look_ahead, 0])  # 对应每个X序列的滞后一期序列值
    return np.array(X, dtype=np.float32), np.array(Y, dtype=np.float32)


def shuffle_data(data, labels):
    """ Shuffle data and labels.
        Input:
          data: B,N,... numpy array
          label: B,... numpy array
        Return:
          shuffled data, label and shuffle indices
    """
    idx = np.arange(len(labels))  # 生成labels个数元素的数组
    np.random.shuffle(idx)  # 打乱数组元素
    return data[idx, ...], labels[idx], idx
